"""
该脚本示例了一个 WebSocket 客户端 + SIM 模拟的使用，
分为两个线程循环
- 生产者用于连接服务器并接收数据放置在队列中, 并定时处理逻辑(双协程)
- 消费者用于消耗队列中的数据, 并进行处理, 如果有必要可以将TX数据await发送给服务器
其次示例代码实现了交互的测试功能,
接收包括
- 验证信息
- 交易数据(Tick)
- 订单信息
- 仓位信息
发送包括
- 订阅Tick数据
- 多头开仓指令
- 多头平仓指令
- 空头开仓指令
- 空头平仓指令
- 撤单指令
- 查询当前仓位
- 查询历史订单
"""

from __future__ import annotations

import asyncio
import dataclasses
import logging
import random
import sys
import traceback

import dotenv
import httpx

from eztrade_client import WebSocketClient
from eztrade_client import (
    Direction,
    Operation,
    Order,
    OrderOperation,
    OrderStatus,
    PyMessage,
    PositionEffect,
    Portfolio,
    TickData,
    get_logger,
)

dotenv.load_dotenv()  # 从 .env 文件加载环境变量


SERVER = "127.0.0.1"
PORT = 7878
BASE_URL = f"https://{SERVER}:{PORT}"
STAGE_ID = 2
PATH = f"/ws/{STAGE_ID}"
USER_ID = 1
STAGE_PAYLOAD = {
    "id": STAGE_ID,
    "stage_type": "SIM",
    "name": f"test-sim-{STAGE_ID}",
    "description": "测试 SIM Stage",
}
SYMBOL = "au2512"  # 主要交易品种

logger = get_logger("Simulate", level=logging.INFO)
logger.addFileHandler("logs/simulate.log", level=logging.DEBUG)


@dataclasses.dataclass
class AccountsInfo:
    sent_orders: dict[int, Order] = dataclasses.field(default_factory=dict)
    done_orders: dict[int, Order] = dataclasses.field(default_factory=dict)
    portfolio: Portfolio | None = dataclasses.field(default=None)
    lastprice: dict[str, float] = dataclasses.field(default_factory=dict)


ACCOUNTS_INFO = AccountsInfo()


async def process_events(client: WebSocketClient) -> None:
    """持续消耗数据队列, 并根据不同类型的事件执行相应操作。"""

    while True:
        if client.buffer:
            data = client.buffer.popleft()
            if len(client.buffer) > 10:
                logger.warning("Size of buffer: %s", len(client.buffer))
            try:
                message = PyMessage.from_binary_bytes(data)
            except Exception as exc:  # noqa: BLE001
                logger.error("Failed to decode message: %s, raw: %s", exc, data)
                continue

            match message.get_data_type():
                case "Authority":
                    auth = message.get_data()
                    logger.debug("UserId: %s", auth.client_id)
                    logger.info("Query current positions")
                    client.send_message(PyMessage.query_account())
                    client.send_message(PyMessage.subscribe_tick(SYMBOL))
                case "TickData":
                    tickdata: TickData = message.get_data()
                    ACCOUNTS_INFO.lastprice[tickdata.symbol] = tickdata.last
                    logger.debug("Received TickData: %s", tickdata)
                    direction = random.choice([Direction.Long, Direction.Short])
                    position_effect = random.choice([
                        PositionEffect.Open,
                        PositionEffect.Close,
                    ])
                    order = Order(tickdata.symbol, direction, position_effect, 1)
                    msg = PyMessage(order)
                    client.send_message(msg)
                    logger.info("Send Order: %s", order)
                case "Order":
                    order: Order = message.get_data()
                    if order.id not in ACCOUNTS_INFO.sent_orders:
                        logger.info("Order %s has been sent.", order.id)
                        ACCOUNTS_INFO.sent_orders.update({order.id: order})
                        if (
                            ACCOUNTS_INFO.portfolio
                            and ACCOUNTS_INFO.portfolio.balance_available / ACCOUNTS_INFO.portfolio.currency.rate()
                            < 5_000_000
                        ):
                            msg = PyMessage(OrderOperation(order.id, Operation.Cancel))
                            logger.debug("Cancel Order: %s", order.id)
                            client.send_message(msg)
                    if order.id in ACCOUNTS_INFO.sent_orders:
                        if order.status in [
                            OrderStatus.Filled,
                            OrderStatus.PartiallyFilled,
                            OrderStatus.Canceled,
                            OrderStatus.Rejected,
                        ]:
                            logger.info("Queried Order has done! %s", order)
                            ACCOUNTS_INFO.sent_orders.pop(order.id)
                            ACCOUNTS_INFO.done_orders.update({order.id: order})
                case "Portfolio":
                    portfolio: Portfolio = message.get_data()
                    logger.info("Got portfolio information: %s", portfolio)
                    logger.info("Current positions: %s", portfolio.positions)
                    ACCOUNTS_INFO.portfolio = portfolio
                case "String":
                    text: str = message.get_data()
                    logger.info("Received string data: %s", text)
                case _:
                    logger.info("Received data: %s", message)
        else:
            await asyncio.sleep(0.001)

        if client.stop_flag.is_set():
            logger.info("客户端接收停止标志已设置, 退出消费者循环")
            break


async def simulate_query(client: WebSocketClient) -> None:
    """模拟查询：发送查询指令, 定时更新仓位信息。"""

    while True:
        if client.stop_flag.is_set():
            logger.info("客户端控制停止标志已设置, 退出控制循环")
            break
        if not ACCOUNTS_INFO.portfolio:
            await asyncio.sleep(1)
            continue

        client.send_message(PyMessage.query_account())
        if ACCOUNTS_INFO.portfolio.positions:
            value_all = 0.0
            logger.info(
                "current balance: %s",
                ACCOUNTS_INFO.portfolio.balance_available / ACCOUNTS_INFO.portfolio.currency.rate(),
            )
            logger.info(
                "current b_value: %s",
                ACCOUNTS_INFO.portfolio.balance / ACCOUNTS_INFO.portfolio.currency.rate(),
            )
            value_all += ACCOUNTS_INFO.portfolio.balance / ACCOUNTS_INFO.portfolio.currency.rate()
            for position in ACCOUNTS_INFO.portfolio.positions:
                last = ACCOUNTS_INFO.lastprice.get(position.symbol, 0)
                logger.info(
                    "Position: %s, value=%s", position, last * position.quantity
                )
                if position.side == Direction.Long:
                    value_all += last * position.quantity
                elif position.side == Direction.Short:
                    value_all -= last * position.quantity
            logger.info("Total portfolio value: %s", value_all)

        await asyncio.sleep(5)


async def my_strategy(client: WebSocketClient) -> None:  # noqa: ARG001
    """占位符：用户可在此实现自定义策略。"""


async def main() -> None:
    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.post("/stages/create", json=STAGE_PAYLOAD)
        assert resp.status_code in [409, 201], f"Create stage failed: {resp.text}"

    client = WebSocketClient()
    client.create(SERVER, PORT, PATH, USER_ID)
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(process_events(client))
            tg.create_task(simulate_query(client))
    except ExceptionGroup as exc_group:
        for exc in exc_group.exceptions:
            logger.error(
                "主循环异常退出: %s: %s", exc, "".join(traceback.format_exception(exc))
            )
    except asyncio.CancelledError:
        logger.info("主循环取消退出")
    finally:
        client.stop_loop_thread()

    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.delete(f"/stages/{STAGE_ID}")
        if resp.status_code != 204:
            logger.error("Delete stage failed: %s", resp.text)
        else:
            logger.info("Stage %s deleted successfully.", STAGE_ID)

    sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())
